import astro.fit
import matplotlib.pyplot as pl

class InteractiveCoFit(object):
    self.help_message = """
    Left click : add a new continuum point
    'd'        : delete the nearest point
    'q'        : quit
    """

    def __new__(self, wa, fl, er, contpoints=None, nbin=8):
        """ Initialise figure, plots and variables.

        Inputs
        ------
        wa:   Wavelengths
        fl:   Fluxes
        er:   One sigma errors
        nbin: Number of pixels to bin arrays in wavelength. Default 8.
        
        Updates
        -------
        self.spec:  Dictionary of wa, fl, er.
        self.contpoints:  Points used to define the continuum.
        self.nbin:  The input nbin value.
        self.markers:  Dictionary of plotting markers .      
        self.connections:  Callback connections.
        self.fig:  The plotting figure instance.
        """
        #setup
        self.spec = dict(wa=wa, fl=fl, er=er)
        self.nbin = nbin
        self.contpoints = []
        if contpoints is not None:
            self.contpoints.extend(contpoints)
        self.markers = dict()
        self.fig = pl.figure()
        self.connections = []
        self.makefig()
        self.updatefig()
        self.modifypoints()    
        pl.show()

    def makefig(self):
        """ Set up the figure and do initial plots. 

        Updates
        -------
        self.markers
        """
        wa,fl,er = [spec[k][0:-1:nbin] for k in 'wa fl er'.split()]
        # axis for spectrum & continuum
        a0 = self.fig.add_axes((0.05,0.1,0.9,0.6))
        a0.set_autoscale_on(0)
        # axis for residuals
        a1 = self.fig.add_axes((0.05,0.75,0.9,0.2),sharex=a0)
        a1.set_autoscale_on(0)
        a1.axhline(0,color='g')
        a1.axhline(1,color='g',alpha=0.5)
        a1.axhline(-1,color='g',alpha=0.5)
        a1.axhline(2,color='g',linestyle='dotted',alpha=0.5)
        a1.axhline(-2,color='g',linestyle='dotted',alpha=0.5)
        m0, = a1.plot([0],[0],'ok',ms=3,alpha=0.3)
        a1.set_ylim(-4,4)
        a0.axhline(0, color='0.7')
        a0.plot(wa,er, color='crimson', alpha=0.8)
        a0.plot(wa,fl, 'b', linestyle='steps-mid')
        m1, = a0.plot([0],[0],'r',alpha=0.7)
        m2, = a0.plot([0],[0],'rx',ms=10,mew=1,picker=5)
        a0.set_xlim(min(wa),max(wa))
        ymin = -np.median(er)*5
        ymax = sorted(fl)[int(len(fl)*0.95)]*1.5
        a0.set_ylim(ymin, ymax)
        self.fig.canvas.draw()
        self.markers.update(contpoints=m2, cont=m1, resid=m0)

    def updatefig(self):
        """ Calculates the new continuum, residuals and updates the plots.

        Updates
        -------
        self.markers
        """
        nbin = self.nbin
        wa,fl,er = (self.spec[key] for key in 'wa','fl','er')
        x,y = zip(*self.contpoints)
        spline = astro.fit.InterpCubicSpline(x,y)
        cont = spline(wa)
        resid = (fl - cont) / er
        self.markers['contpoints'].set_data(zip(*self.contpoints))
        self.markers['cont'].set_data(wa[0:-1:nbin], cont[0:-1:nbin])
        self.markers['resid'].set_data(wa[0:-1:nbin], resid[0:-1:nbin])
        self.fig.canvas.draw()

    def on_click(self, event):
        """ Add a new continuum point.

        Updates
        -------
        self.contpoints
        """
        if self.fig.canvas.toolbar._active in ('ZOOM','PAN'):  return
        if event.inaxes != self.fig.axes[0]:  return
        # add a point to contpoints
        x,y = event.xdata,event.ydata
        if x not in zip(*self.contpoints)[0]:
            self.contpoints.append((x,y))
            self.updatefig()

    def on_keypress(self, event):
        """ Remove a continuum point.

        Updates
        -------
        self.contpoints
        """
        if self.fig.canvas.toolbar._active in ('ZOOM','PAN'):
            print 'Must not be in zoom/pan mode'
            return
        if event.inaxes != self.fig.axes[0]:  return
        if event.key == 'd':
            # remove a point from contpoints
            contx,conty = zip(*self.contpoints)
            sep = np.hypot(event.xdata - contx, event.ydata - conty)
            self.contpoints.remove((self.contpoints[sep.argmin()]))
            self.updatefig()
        elif event.key == '?':
            print self.help_message
        elif event.key == 'q':
            for item in self.connections:
                self.fig.canvas.mpl_disconnect(item)

    def modifypoints(self):
        """ Add/remove continuum points."""
        print self.help_message
        id0 = self.fig.canvas.mpl_connect('key_press_event',self.on_keypress)
        id1 = self.fig.canvas.mpl_connect('button_press_event',self.on_click)
        self.connections.extend([id0,id1])
